#include <cassert>
#include <utility>
#include <cstdlib>
#include "dft_inplace/dft.h"
#include "dft.h"

using namespace std;
using namespace tbb;

namespace tbbimpl {

  class ButterflyCalculator {
  public:
    ButterflyCalculator(complex<float>* _x) : x(_x) {}

    void operator ()(const blocked_range<size_t>& range) const {
      for (size_t i = range.begin(); i < range.end(); i += 2*size) {
        for (int j = 0; j < size; ++j) {
          complex<float> a = x[i + j];
          complex<float> b = x[i + j + size];
          float ang_ = ang * j;
          b *= complex<float>(cos(ang_), sin(ang_));
          x[i + j] = a + b;
          x[i + j + size] = a - b;
        }
      }
    }

    int size;
    float ang;
    complex<float> *x;
  };

  void DftInplace(complex<float>* x, int n) {
    assert(IsPowerOfTwo(n));
    ReorderElements(x, n);

    ButterflyCalculator calc(x); 
    for (int size = 1; size < n; size <<= 1) {
      calc.size = size;
      calc.ang = -static_cast<float>(M_PI) / size;
      parallel_for(blocked_range<size_t>(0, n, 2 * size), calc);
    }
  }

} // namespace tbbimpl